/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.gaia.translation.functions;



import com.bff.gaia.unified.runners.core.SideInputReader;

import com.bff.gaia.unified.sdk.transforms.Materializations;

import com.bff.gaia.unified.sdk.transforms.Materializations.MultimapView;

import com.bff.gaia.unified.sdk.transforms.ViewFn;

import com.bff.gaia.unified.sdk.transforms.windowing.BoundedWindow;

import com.bff.gaia.unified.sdk.values.PCollectionView;

import com.bff.gaia.unified.sdk.values.TupleTag;

import com.bff.gaia.unified.sdk.values.WindowingStrategy;

import com.bff.gaia.api.common.functions.RuntimeContext;



import javax.annotation.Nullable;

import java.util.Collections;

import java.util.HashMap;

import java.util.Map;



import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions.checkArgument;

import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions.checkNotNull;



/** A {@link SideInputReader} for the Gaia Batch Runner. */

public class GaiaSideInputReader implements SideInputReader {

  /** A {@link MultimapView} which always returns an empty iterable. */

  private static final MultimapView EMPTY_MULTMAP_VIEW = o -> Collections.EMPTY_LIST;



  private final Map<TupleTag<?>, WindowingStrategy<?, ?>> sideInputs;



  private RuntimeContext runtimeContext;



  public GaiaSideInputReader(

      Map<PCollectionView<?>, WindowingStrategy<?, ?>> indexByView, RuntimeContext runtimeContext) {

    for (PCollectionView<?> view : indexByView.keySet()) {

      checkArgument(

          Materializations.MULTIMAP_MATERIALIZATION_URN.equals(

              view.getViewFn().getMaterialization().getUrn()),

          "This handler is only capable of dealing with %s materializations "

              + "but was asked to handle %s for PCollectionView with tag %s.",

          Materializations.MULTIMAP_MATERIALIZATION_URN,

          view.getViewFn().getMaterialization().getUrn(),

          view.getTagInternal().getId());

    }

    sideInputs = new HashMap<>();

    for (Map.Entry<PCollectionView<?>, WindowingStrategy<?, ?>> entry : indexByView.entrySet()) {

      sideInputs.put(entry.getKey().getTagInternal(), entry.getValue());

    }

    this.runtimeContext = runtimeContext;

  }



  @Nullable

  @Override

  public <T> T get(PCollectionView<T> view, BoundedWindow window) {

    checkNotNull(view, "View passed to sideInput cannot be null");

    TupleTag<?> tag = view.getTagInternal();

    checkNotNull(sideInputs.get(tag), "Side input for " + view + " not available.");



    Map<BoundedWindow, T> sideInputs =

        runtimeContext.getBroadcastVariableWithInitializer(

            tag.getId(), new SideInputInitializer<>(view));

    T result = sideInputs.get(window);

    if (result == null) {

      ViewFn<MultimapView, T> viewFn = (ViewFn<MultimapView, T>) view.getViewFn();

      result = viewFn.apply(EMPTY_MULTMAP_VIEW);

    }

    return result;

  }



  @Override

  public <T> boolean contains(PCollectionView<T> view) {

    return sideInputs.containsKey(view.getTagInternal());

  }



  @Override

  public boolean isEmpty() {

    return sideInputs.isEmpty();

  }

}